// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/exec/window_function.h>
#include <pollux/expression/function_signature.h>
#include <pollux/expression/signature_binder.h>

namespace kumo::pollux::exec {
    WindowFunctionMap &windowFunctions() {
        static WindowFunctionMap functions;
        return functions;
    }

    namespace {
        std::optional<const WindowFunctionEntry *> getWindowFunctionEntry(
            const std::string &name) {
            auto &functionsMap = windowFunctions();
            auto it = functionsMap.find(name);
            if (it != functionsMap.end()) {
                return &it->second;
            }

            return std::nullopt;
        }
    } // namespace

    bool registerWindowFunction(
        const std::string &name,
        std::vector<FunctionSignaturePtr> signatures,
        WindowFunction::Metadata metadata,
        WindowFunctionFactory factory) {
        auto sanitizedName = sanitize_name(name);
        windowFunctions()[sanitizedName] = {
            std::move(signatures), std::move(factory), std::move(metadata)
        };
        return true;
    }

    WindowFunction::Metadata getWindowFunctionMetadata(const std::string &name) {
        const auto sanitizedName = sanitize_name(name);
        if (auto func = getWindowFunctionEntry(sanitizedName)) {
            return func.value()->metadata;
        } else {
            POLLUX_USER_FAIL(
                "Window function metadata not found for function: {}", name);
        }
    }

    std::optional<std::vector<FunctionSignaturePtr> > getWindowFunctionSignatures(
        const std::string &name) {
        auto sanitizedName = sanitize_name(name);
        if (auto func = getWindowFunctionEntry(sanitizedName)) {
            return func.value()->signatures;
        }
        return std::nullopt;
    }

    std::unique_ptr<WindowFunction> WindowFunction::create(
        const std::string &name,
        const std::vector<WindowFunctionArg> &args,
        const TypePtr &resultType,
        bool ignoreNulls,
        memory::MemoryPool *pool,
        HashStringAllocator *stringAllocator,
        const core::QueryConfig &config) {
        // Lookup the function in the new registry first.
        if (auto func = getWindowFunctionEntry(name)) {
            std::vector<TypePtr> argTypes;
            argTypes.reserve(args.size());
            for (const auto &arg: args) {
                argTypes.push_back(arg.type);
            }

            const auto &signatures = func.value()->signatures;
            for (auto &signature: signatures) {
                SignatureBinder binder(*signature, argTypes);
                if (binder.tryBind()) {
                    auto type = binder.tryResolveType(signature->returnType());
                    POLLUX_USER_CHECK(
                        type->equivalent(*resultType),
                        "Unexpected return type for window function {}. Expected {}. Got {}.",
                        toString(name, argTypes),
                        type->toString(),
                        resultType->toString());
                    return func.value()->factory(
                        args, resultType, ignoreNulls, pool, stringAllocator, config);
                }
            }

            POLLUX_USER_FAIL(
                "Window function signature is not supported: {}. Supported signatures: {}.",
                toString(name, argTypes),
                toString(signatures));
        }

        POLLUX_USER_FAIL("Window function not registered: {}", name);
    }

    void WindowFunction::setEmptyFramesResult(
        const SelectivityVector &validRows,
        vector_size_t resultOffset,
        const VectorPtr &defaultResult,
        const VectorPtr &result) {
        if (validRows.isAllSelected()) {
            return;
        }
        // Set the null bit for all rows to true if the defaultResult is NULL and
        // there are no valid rows.
        if (!validRows.hasSelections() && defaultResult->is_null_at(0)) {
            uint64_t *rawNulls = result->mutable_raw_nulls();
            bits::fillBits(
                rawNulls, resultOffset, resultOffset + validRows.size(), bits::kNull);
            return;
        }

        invalidRows_.resizeFill(validRows.size(), true);
        invalidRows_.deselect(validRows);
        invalidRows_.applyToSelected([&](auto i) {
            result->copy(defaultResult.get(), resultOffset + i, 0, 1);
        });
    }
} // namespace kumo::pollux::exec
